import os
import sys
import binascii
import struct
import zlib
import itertools
import re
from collections import defaultdict
from PIL import Image
from PySide6.QtWidgets import (
    QMainWindow,
    QWidget,
    QVBoxLayout,
    QLabel,
    QLineEdit,
    QPushButton,
    QFileDialog,
    QTextEdit,
    QToolBar,
    QApplication,
)
from PySide6.QtGui import QAction, QIcon
import ctypes


def resource_path(relative_path):
    """获取资源文件的绝对路径（适用于打包后的程序）"""
    try:
        # PyInstaller 创建的临时文件夹
        base_path = sys._MEIPASS
    except Exception:
        base_path = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(base_path, relative_path)


# 定义文件签名（Magic Numbers）
SIGNATURES = [
    (b'\x89\x50\x4E\x47\x0D\x0A\x1A\x0A', 'PNG', 'png'),
    (b'\x47\x49\x46\x38\x37\x61', 'GIF', 'gif'),
    (b'\x47\x49\x46\x38\x39\x61', 'GIF', 'gif'),
    (b'\xFF\xD8\xFF', 'JPEG', 'jpg'),
    (b'\x25\x50\x44\x46\x2D', 'PDF', 'pdf'),
    (b'\x50\x4B\x03\x04', 'ZIP', 'zip'),
    (b'\x1F\x8B\x08', 'GZIP', 'gz'),
    (b'\x52\x61\x72\x21\x1A\x07\x00', 'RAR', 'rar'),
    (b'\x52\x61\x72\x21\x1A\x07\x01\x00', 'RAR5', 'rar'),
    (b'\x4D\x5A\x90\x00', 'PE32', 'exe'),
]


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle("图像处理工具")
        self.setGeometry(100, 100, 600, 450)

        # 文件路径输入框
        self.file_path_label = QLabel("图片路径:", self)
        self.file_path_input = QLineEdit(self)
        self.file_path_input.setReadOnly(True)

        # 文件选择按钮
        self.select_file_button = QPushButton("选择文件", self)
        self.select_file_button.clicked.connect(self.select_file)

        # 一把梭按钮
        self.one_click_button = QPushButton("一把梭", self)
        self.one_click_button.clicked.connect(self.one_click_process)

        # 输出日志
        self.log_text = QTextEdit(self)
        self.log_text.setReadOnly(True)

        # 布局
        layout = QVBoxLayout()
        layout.addWidget(self.file_path_label)
        layout.addWidget(self.file_path_input)
        layout.addWidget(self.select_file_button)
        layout.addWidget(self.one_click_button)
        layout.addWidget(QLabel("执行日志:"))
        layout.addWidget(self.log_text)

        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

        # 添加工具栏
        self.add_toolbars()

        # 设置背景图片和图标
        self.set_background_and_icon()

    def set_background_and_icon(self):
        # 设置背景图片
        background_path = resource_path("img/OIP-C.jpg(1)")
        self.setStyleSheet(f"MainWindow {{ background-image: url({background_path}); background-size: cover; }}")

        # 设置窗口图标
        icon_path = resource_path("img/OIP-C.jpg")
        self.setWindowIcon(QIcon(icon_path))

    def add_toolbars(self):
        # 添加工具栏
        toolbar = QToolBar("工具栏")
        self.addToolBar(toolbar)

        # 添加动作
        action1 = QAction("清屏", self)
        action1.triggered.connect(self.clear_log)
        toolbar.addAction(action1)

    def select_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片文件", "", "图片文件 (*.png *.jpg *.jpeg *.gif *.bmp)"
        )
        if file_path:
            self.file_path_input.setText(file_path)
            self.log(f"已选择文件: {file_path}")
            self.add_separator()

    def clear_log(self):
        self.log_text.clear()
        with open("result.txt", "w") as f:
            f.write("")

    def log(self, message):
        self.log_text.append(message)
        with open("result.txt", "a", encoding='utf-8') as f:
            f.write(message + "\n")

    def add_separator(self):
        separator = "-" * 50
        self.log(separator)

    def image_to_hex(self):
        file_path = self.file_path_input.text()
        if not file_path:
            self.log("请先选择图片文件！")
            return

        try:
            with open(file_path, 'rb') as image_file:
                binary_data = image_file.read()

            hex_data = binascii.hexlify(binary_data).decode('utf-8')

            with open("hex.txt", 'w') as hex_file:
                hex_file.write(hex_data)
            self.log("图片已成功转换为16进制，保存到 hex.txt")
            self.add_separator()

        except Exception as e:
            self.log(f"转换图片为16进制时出错: {e}")
            self.add_separator()

    def process_image(self):
        file_path = self.file_path_input.text()
        if not file_path:
            self.log("请先选择图片文件！")
            return

        try:
            # 检查图片格式是否为 PNG
            with Image.open(file_path) as img:
                if img.format != 'PNG':
                    self.log("图片不是 PNG 格式，跳过宽高爆破。")
                    self.add_separator()
                    return

            with open(file_path, 'rb') as file:
                bin_data = file.read()

            crc32key = zlib.crc32(bin_data[12:29])
            original_crc32 = int(bin_data[29:33].hex(), 16)

            if crc32key == original_crc32:
                self.log("宽高没有问题！")
                self.add_separator()
            else:
                self.log("开始爆破宽高...")
                for i, j in itertools.product(range(4096), range(4096)):
                    data = bin_data[12:16] + struct.pack('>i', i) + struct.pack('>i', j) + bin_data[24:29]
                    crc32 = zlib.crc32(data)
                    if crc32 == original_crc32:
                        self.log(f"找到正确宽高: 宽度={i}, 高度={j}")
                        self.modify_hex_file(i, j)
                        return
                self.log("破解失败")
                self.add_separator()

        except Exception as e:
            self.log(f"处理图片时出错: {e}")
            self.add_separator()

    def modify_hex_file(self, width, height):
        try:
            with open('hex.txt', 'r') as hex_file:
                hex_data = hex_file.read()

            bin_data = binascii.unhexlify(hex_data)
            modified_data = bin_data[:16] + struct.pack('>I', width) + struct.pack('>I', height) + bin_data[24:]

            modified_hex_data = binascii.hexlify(modified_data).decode('utf-8')

            with open('hex.txt', 'w') as hex_file:
                hex_file.write(modified_hex_data)

            with open('modified_image.png', 'wb') as png_file:
                png_file.write(modified_data)

            self.log("修改后的16进制数据已保存到 hex.txt，修改后的图片已保存为 modified_image.png")
            self.add_separator()

        except Exception as e:
            self.log(f"修改 hex.txt 文件时出错: {e}")
            self.add_separator()

    def scan_hidden_files(self):
        input_file = 'hex.txt'
        if not os.path.exists(input_file):
            self.log(f"Error: Input file '{input_file}' not found.")
            self.add_separator()
            return

        try:
            with open(input_file, 'r', encoding='utf-8') as f:
                hex_content = f.read()

            hex_data = re.findall(r'[0-9A-Fa-f]{2}', hex_content)
            hex_str = ''.join(hex_data)
            data = binascii.unhexlify(hex_str)

            matches = []
            for signature, file_type, ext in SIGNATURES:
                offset = 0
                while True:
                    offset = data.find(signature, offset)
                    if offset == -1:
                        break
                    matches.append((offset, signature, file_type, ext))
                    offset += 1

            if not matches:
                self.log("No files found in the hex data.")
                self.add_separator()
                return

            self.log(f"Found {len(matches)} potential files:")
            for offset, _, file_type, _ in matches:
                self.log(f"  {file_type} at offset 0x{offset:X}")

            # 提取文件
            output_dir = 'extracted_files'
            self.log("文件已经成功提取到extracted_files文件夹中")
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)

            for idx, (offset, signature, file_type, ext) in enumerate(matches):
                file_data = data[offset:]
                file_name = f"{file_type}_{idx}.{ext}"
                file_path = os.path.join(output_dir, file_name)

                with open(file_path, 'wb') as f:
                    f.write(file_data)

                self.log(f"Extracted {file_type} file: {file_name} (Offset: 0x{offset:X})")
            self.add_separator()

        except Exception as e:
            self.log(f"Error scanning hidden files: {e}")
            self.add_separator()

    def get_image_info(self):
        file_path = self.file_path_input.text()
        if not file_path:
            self.log("请先选择图片文件！")
            return

        try:
            with Image.open(file_path) as img:
                self.log("图片属性详细信息：")
                self.log(f"文件路径: {file_path}")
                self.log(f"文件格式: {img.format}")
                self.log(f"图片尺寸: {img.size} (宽度 x 高度)")
                self.log(f"图片模式: {img.mode}")

                exif_data = img._getexif()
                if exif_data:
                    self.log("\n图片属性详细信息：")
                    from PIL.ExifTags import TAGS
                    for tag, value in exif_data.items():
                        tag_name = TAGS.get(tag, tag)
                        self.log(f"{tag_name}: {value}")
                else:
                    self.log("\n该图片没有属性详细信息。")
            self.add_separator()

        except Exception as e:
            self.log(f"处理图片时出错: {e}")
            self.add_separator()

    def string_search(self):
        input_file = 'hex.txt'
        if not os.path.exists(input_file):
            self.log(f"Error: Input file '{input_file}' not found.")
            self.add_separator()
            return

        try:
            with open(input_file, 'r', encoding='utf-8') as f:
                hex_content = f.read()

            # 将16进制字符串转换为字节流
            hex_data = re.findall(r'[0-9A-Fa-f]{2}', hex_content)
            hex_str = ''.join(hex_data)
            data = binascii.unhexlify(hex_str)

            # 将字节流解码为字符串（忽略非ASCII字符）
            text = data.decode('utf-8', errors='replace')

            # 定义要搜索的关键词
            keywords = ["ctf", "CTF", "FLAG", "flag", "{}", "ZmxhZwo", "Y3RmCg", "password", "pd", "=="]

            # 搜索文件中的关键词
            matches = []
            for keyword in keywords:
                index = text.find(keyword)
                while index != -1:
                    # 获取匹配字符串的上下文
                    start = max(0, index - 32)
                    end = min(len(text), index + len(keyword) + 32)
                    context = text[start:end]
                    matches.append((keyword, context))
                    index = text.find(keyword, index + 1)

            if not matches:
                self.log("未找到关键词。")
            else:
                self.log(f"找到 {len(matches)} 个关键词:")
                for keyword, context in matches:
                    self.log(f"  关键词: {keyword}")
                    self.log(f"  上下文: {context}\n")
            self.add_separator()

        except Exception as e:
            self.log(f"字符串搜索时出错: {e}")
            self.add_separator()

    def lsb_decrypt(self):
        file_path = self.file_path_input.text()
        if not file_path:
            self.log("请先选择图片文件！")
            return

        try:
            # 打开图片
            with Image.open(file_path) as img:
                if img.mode not in ['RGB', 'RGBA']:
                    self.log("图片模式不支持 LSB 隐写解密，请使用 RGB 或 RGBA 模式的图片。")
                    self.add_separator()
                    return

                # 将图片转换为 RGB 模式
                img = img.convert('RGB')
                pixels = list(img.getdata())

                # 提取 LSB
                binary_data = []
                for pixel in pixels:
                    for channel in pixel:
                        # 提取最低有效位
                        binary_data.append(channel & 1)

                # 将二进制数据转换为字符串
                binary_str = ''.join(map(str, binary_data))

                # 将二进制数据转换为16进制
                hex_data = []
                for i in range(0, len(binary_str), 8):
                    byte = binary_str[i:i + 8]
                    if len(byte) < 8:
                        break
                    hex_byte = hex(int(byte, 2))[2:].zfill(2)
                    hex_data.append(hex_byte)

                # 将16进制数据转换为字符串
                hex_str = ''.join(hex_data)

                # 将16进制数据保存到 lsb_hex.txt
                with open("lsb_hex.txt", "w") as f:
                    f.write(hex_str)
                # 将16进制转化为ASCII码
                try:
                    asc_str = bytes.fromhex(hex_str).decode('utf-8', errors='ignore')
                except ValueError:
                    # 如果转换失败，可能是因为16进制字符串长度不是偶数
                    # 可以尝试修复
                    hex_str = hex_str[:len(hex_str) // 2 * 2]  # 截取偶数长度
                    asc_str = bytes.fromhex(hex_str).decode('utf-8', errors='ignore')
                # 在日志中显示前30位数据
                self.log("LSB 隐写解密结果:")
                self.log(f"前30位ascII数据: {asc_str[:40]}")
                self.log(f"16进制数据已保存到 lsb_hex.txt")
                self.add_separator()

        except Exception as e:
            self.log(f"LSB 隐写解密时出错: {e}")
            self.add_separator()

    def scan_lsb_hex_files(self):
        input_file = 'lsb_hex.txt'
        if not os.path.exists(input_file):
            self.log(f"Error: Input file '{input_file}' not found.")
            self.add_separator()
            return

        try:
            with open(input_file, 'r', encoding='utf-8') as f:
                hex_content = f.read()

            # 将16进制字符串转换为字节流
            hex_data = re.findall(r'[0-9A-Fa-f]{2}', hex_content)
            hex_str = ''.join(hex_data)
            data = binascii.unhexlify(hex_str)

            matches = []
            for signature, file_type, ext in SIGNATURES:
                offset = 0
                while True:
                    offset = data.find(signature, offset)
                    if offset == -1:
                        break
                    matches.append((offset, signature, file_type, ext))
                    offset += 1

            if not matches:
                self.log("No files found in the LSB hex data.")
                self.add_separator()
                return

            self.log(f"Found {len(matches)} potential files in LSB hex data:")
            for offset, _, file_type, _ in matches:
                self.log(f"  {file_type} at offset 0x{offset:X}")

            # 提取文件
            output_dir = 'lsb_extracted_files'
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)

            for idx, (offset, signature, file_type, ext) in enumerate(matches):
                file_data = data[offset:]
                file_name = f"LSB_{file_type}_{idx}.{ext}"
                file_path = os.path.join(output_dir, file_name)

                with open(file_path, 'wb') as f:
                    f.write(file_data)

                self.log(f"Extracted {file_type} file from LSB hex data: {file_name} (Offset: 0x{offset:X})")
            self.add_separator()

        except Exception as e:
            self.log(f"Error scanning LSB hex data for hidden files: {e}")
            self.add_separator()

    def extract_channels(self):
        file_path = self.file_path_input.text()
        if not file_path:
            self.log("请先选择图片文件！")
            return

        try:
            with Image.open(file_path) as img:
                if img.mode not in ['RGB', 'RGBA']:
                    self.log("图片模式不支持通道提取，请使用 RGB 或 RGBA 模式的图片。")
                    self.add_separator()
                    return

                # 分离通道
                if img.mode == 'RGB':
                    r, g, b = img.split()
                    channels = {'R': r, 'G': g, 'B': b}
                else:  # RGBA
                    r, g, b, a = img.split()
                    channels = {'R': r, 'G': g, 'B': b, 'A': a}

                # 保存原始通道
                output_dir = 'extracted_channels'
                if not os.path.exists(output_dir):
                    os.makedirs(output_dir)

                for channel_name, channel_img in channels.items():
                    channel_path = os.path.join(output_dir, f"{channel_name}_channel.png")
                    channel_img.save(channel_path)
                    self.log(f"通道 {channel_name} 已保存为 {channel_path}")

                # 提取位平面
                for channel_name, channel_img in channels.items():
                    # 创建位平面文件夹
                    bit_plane_dir = os.path.join(output_dir, f"{channel_name}_bit_planes")
                    if not os.path.exists(bit_plane_dir):
                        os.makedirs(bit_plane_dir)

                    # 提取每个位平面（0到7）
                    for bit in range(8):
                        # 创建一个新的图像来存储当前位平面
                        bit_plane_img = Image.new('L', channel_img.size)
                        pixels = bit_plane_img.load()

                        # 提取当前位平面
                        for y in range(channel_img.size[1]):
                            for x in range(channel_img.size[0]):
                                pixel = channel_img.getpixel((x, y))
                                # 提取第 bit 位，并将其放大为 0 或 255 以提高可见性
                                pixels[x, y] = (pixel >> bit) & 1
                                # 将位平面转换为黑白图像（0 或 255）
                                pixels[x, y] = pixels[x, y] * 255

                        # 保存位平面
                        bit_plane_path = os.path.join(bit_plane_dir, f"{channel_name}_bit_plane_{bit}.png")
                        bit_plane_img.save(bit_plane_path)
                self.log(f"成功保存在extracted_channels")
                self.add_separator()

        except Exception as e:
            self.log(f"提取通道或位平面时出错: {e}")
            self.add_separator()

    def one_click_process(self):
        self.log("开始执行'一把梭'流程...")
        self.add_separator()
        self.image_to_hex()
        self.process_image()
        self.scan_hidden_files()
        self.get_image_info()
        self.string_search()
        self.lsb_decrypt()
        self.scan_lsb_hex_files()
        self.extract_channels()
        self.log("流程执行完成！")
        self.add_separator()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())